import random

from common.random_util import select_one_from_weight_map


class SrWeaponPool:
    ABSTRACT_5 = 0
    WEAPON_5_NORMAL = 1
    WEAPON_5_UP = 3
    ABSTRACT_4 = 4
    ABSTRACT_3 = 5

    WEIGHT_LIMIT = 10000

    def __init__(self):
        self.not_5_count = 0
        self.not_4_count = 0
        self.must_up = False

    def next(self):
        res = select_one_from_weight_map(self._build_weight_map())
        if res == SrWeaponPool.ABSTRACT_5:
            self.not_5_count = 0
            if self.must_up:
                self.must_up = False
                return SrWeaponPool.WEAPON_5_UP
            if random.random() < 0.75:
                self.must_up = False
                return SrWeaponPool.WEAPON_5_UP
            else:
                self.must_up = True
                return SrWeaponPool.WEAPON_5_NORMAL
        self.not_5_count += 1
        if res == SrWeaponPool.ABSTRACT_4:
            self.not_4_count = 0
        else:
            self.not_4_count += 1
        return res

    def _build_weight_map(self):
        weight_map = {}
        i = self.not_5_count + 1
        j = self.not_4_count + 1

        if i < 63:
            weight_5 = 70
        elif 63 <= i <= 73:
            weight_5 = 70 + 700 * (i - 62)
        else:
            weight_5 = 7700 + 350 * (i - 73)
        weight_map[SrWeaponPool.ABSTRACT_5] = weight_5
        if weight_5 >= SrWeaponPool.WEIGHT_LIMIT:
            return weight_map

        if j <= 7:
            weight_4 = 600
        elif j == 8:
            weight_4 = 6600
        else:
            weight_4 = 6600 + 3000 * (j - 8)
        weight_4 = min(weight_4, SrWeaponPool.WEIGHT_LIMIT - weight_5)
        weight_map[SrWeaponPool.ABSTRACT_4] = weight_4
        if (weight_4 + weight_5) >= SrWeaponPool.WEIGHT_LIMIT:
            return weight_map

        weight_3 = min(9330, SrWeaponPool.WEIGHT_LIMIT - weight_5 - weight_4)
        weight_map[SrWeaponPool.ABSTRACT_3] = weight_3

        return weight_map
